import os
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import mne
import glob
import re

plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'SimSun', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

def ensure_dir(path: str) -> None:
    if not os.path.exists(path):
        os.makedirs(path, exist_ok=True)

def load_raw(vhdr_path: str) -> mne.io.BaseRaw:
    raw = mne.io.read_raw_brainvision(vhdr_path, preload=True, verbose=False)
    '''
    if 'TP10' in raw.ch_names:
        raw.set_eeg_reference(ref_channels=['TP10'], verbose=False)
        print('参考电极: TP10')
    else:
        raw.set_eeg_reference('average', verbose=False)
        print('参考电极: 平均参考 (TP10 不存在)')
        
        '''
    has_tp9 = 'TP9' in raw.ch_names
    has_tp10 = 'TP10' in raw.ch_names
    if has_tp9 and has_tp10:
        raw.set_eeg_reference(ref_channels=['TP9', 'TP10'], verbose=False)
        print('参考电极: TP9 与 TP10 的平均')
    elif has_tp10:
        raw.set_eeg_reference(ref_channels=['TP10'], verbose=False)
        print('参考电极: TP10')
    elif has_tp9:
        raw.set_eeg_reference(ref_channels=['TP9'], verbose=False)
        print('参考电极: TP9')
    else:
        raw.set_eeg_reference('average', verbose=False)
        print('参考电极: 平均参考 (TP9/TP10 不存在)')
        
    try:
        raw.notch_filter(50.0, verbose=False)
    except Exception:
        pass
    raw.filter(3.0, 45.0, fir_design='firwin', verbose=False)
    return raw

def pick_occipital_names(raw: mne.io.BaseRaw) -> list:
    # computer_one
    preferred = ['O1', 'O2', 'Oz', 'P4', 'P8','P7','Pz']
    # AR
    return [ch for ch in preferred if ch in raw.ch_names]

def next_pow2(n: int) -> int:
    return 1 << (int(np.ceil(np.log2(max(1, n)))))

def compute_windowed_amplitude_spectrum(signal_uV: np.ndarray, fs: float, 
                                      window_length: float = 2.0, 
                                      overlap: float = 0.5):
    """
    使用多个窗口叠加计算幅度谱
    
    Parameters:
    -----------
    signal_uV : np.ndarray
        输入信号（µV）
    fs : float
        采样频率（Hz）
    window_length : float
        窗口长度（秒）
    overlap : float
        窗口重叠比例（0-1）
    
    Returns:
    --------
    freqs : np.ndarray
        频率轴
    amp_mean : np.ndarray
        平均幅度谱
    amp_std : np.ndarray
        幅度谱标准差
    """
    # 计算窗口样本数
    window_samples = int(window_length * fs)
    
    # 计算重叠样本数
    overlap_samples = int(window_samples * overlap)
    step_samples = window_samples - overlap_samples
    
    # 创建汉明窗
    window = np.hamming(window_samples)
    
    # 计算FFT点数（使用2的幂次）
    NFFT = next_pow2(window_samples)
    
    # 存储所有窗口的幅度谱
    all_spectra = []
    
    # 分段处理
    start = 0
    while start + window_samples <= len(signal_uV):
        # 提取当前窗口的数据
        segment = signal_uV[start:start + window_samples]
        
        # 去直流
        segment = segment - np.mean(segment)
        
        # 加窗
        segment_windowed = segment * window
        
        # FFT
        Y = np.fft.rfft(segment_windowed, n=NFFT) / window_samples
        amp = np.abs(Y)
        
        # 单边谱修正
        if len(amp) > 2:
            amp[1:-1] *= 2
        
        all_spectra.append(amp)
        
        # 移动到下一个窗口
        start += step_samples
    
    if not all_spectra:
        raise ValueError("信号长度不足以创建窗口")
    
    # 转换为numpy数组
    all_spectra = np.array(all_spectra)
    
    # 计算频率轴
    freqs = np.linspace(0, fs/2, len(all_spectra[0]))
    
    # 计算统计量
    amp_mean = np.mean(all_spectra, axis=0)
    amp_std = np.std(all_spectra, axis=0)
    
    return freqs, amp_mean, amp_std, len(all_spectra)

def annotate_harmonics(ax, freqs, amp, f0: float, max_h: int = 4, search_bw: float = 0.6):
    """标注谐波"""
    def find_peak(f_center: float):
        mask = (freqs >= max(0, f_center - search_bw)) & (freqs <= f_center + search_bw)
        if not np.any(mask):
            return None
        idx = np.argmax(amp[mask])
        idxs = np.where(mask)[0]
        ii = idxs[idx]
        return freqs[ii], amp[ii]

    found = {}
    for k in range(1, max_h + 1):
        fk = f0 * k
        if fk > freqs[-1]:
            break
        ax.axvline(fk, color='gray', linestyle=':', linewidth=1.0)
        peak = find_peak(fk)
        if peak is None:
            continue
        fx, fy = peak
        found[k] = (fx, fy)
        ax.scatter([fx], [fy], marker='s', s=36, color='black', zorder=5)
        ax.annotate(
            f"{k}f\nX: {fx:.2f}\nY: {fy:.5g}",
            xy=(fx, fy), xytext=(5, 8), textcoords='offset points',
            bbox=dict(boxstyle='round,pad=0.3', fc='wheat', alpha=0.8),
            arrowprops=dict(arrowstyle='-'), fontsize=9
        )
    return found

def process_single_file(vhdr_path: str, out_dir: str, window_length: float = 4.0, overlap: float = 0.5):
    """处理单个vhdr文件"""
    try:
        # 从文件名提取基本信息
        base_name = os.path.splitext(os.path.basename(vhdr_path))[0]
        
        # 从文件名提取目标频率
        freq_match = re.search(r'_(\d+(?:\.\d+)?)\.vhdr$', vhdr_path)
        if freq_match:
            target_f = float(freq_match.group(1))
        else:
            print(f"警告：无法从文件名 {vhdr_path} 提取目标频率，跳过此文件")
            return False
        
        print(f"\n正在处理文件: {base_name} (目标频率: {target_f} Hz)")
        
        # 加载数据
        raw = load_raw(vhdr_path)
        fs = raw.info['sfreq']
        picks = pick_occipital_names(raw)
        if len(picks) == 0:
            picks = raw.ch_names
        data = raw.get_data(picks=picks)

        # 通道平均
        signal_uV = (data.mean(axis=0)) * 1e6

        # 提取第5-15秒的数据
        start_sample = int(5 * fs)  # 第5秒开始
        end_sample = int(15 * fs)   # 第15秒结束
        signal_segment = signal_uV[start_sample:end_sample]
        
        print(f"分析时间段: 第5-15秒 ({len(signal_segment)/fs:.1f}秒数据)")

        # 计算窗口化幅度谱
        freqs, amp_mean, amp_std, num_windows = compute_windowed_amplitude_spectrum(
            signal_segment, fs, window_length, overlap
        )

        # 绘图
        fig, ax1 = plt.subplots(1, 1, figsize=(12, 6))

        # 平均幅度谱
        ax1.plot(freqs, amp_mean, color='#32cd32', linewidth=1.5, label='平均幅度谱')
        ax1.fill_between(freqs, amp_mean - amp_std, amp_mean + amp_std, 
                         alpha=0.3, color='#32cd32', label='±1标准差')
        ax1.set_xlim(3, min(45, freqs[-1]))
        ax1.grid(True, which='both', alpha=0.25)
        ax1.set_xlabel('频率 (Hz)')
        ax1.set_ylabel('幅度谱 (\u03bcV)')
        ax1.set_title(f'窗口化幅度谱 (5-15s, 汉明窗, {window_length}s, {overlap*100}%重叠, {num_windows}个窗口)')
        ax1.axvline(target_f, color='crimson', linestyle='--', linewidth=1.5, label=f'{target_f:.2f} Hz')
        ax1.axvline(11.0, color='yellow', linestyle='--', linewidth=1.5, label='11.0 Hz')
        ax1.legend()

        # 标注谐波
        found = annotate_harmonics(ax1, freqs, amp_mean, target_f, max_h=4, search_bw=0.6)

        plt.tight_layout()
        
        # 保存图片
        fig_path = os.path.join(out_dir, f'{base_name}_windowed_amp_5_15s.png')
        plt.savefig(fig_path, dpi=150)
        plt.close(fig)

        # 保存数据
        csv_path = os.path.join(out_dir, f'{base_name}_windowed_amp_5_15s.csv')
        np.savetxt(csv_path, np.c_[freqs, amp_mean, amp_std], delimiter=',', 
                  header='freq(Hz),amplitude_mean(uV),amplitude_std(uV)', comments='')

        # 打印结果
        print("=== 窗口化SSVEP分析结果 (5-15秒) ===")
        print(f"分析时间段: 第5-15秒 ({len(signal_segment)/fs:.1f}秒数据)")
        print(f"窗口长度: {window_length} 秒")
        print(f"重叠比例: {overlap*100}%")
        print(f"窗口数量: {num_windows}")
        print(f"频率分辨率: {fs/next_pow2(int(window_length*fs)):.3f} Hz")
        
        # 目标频率分析
        if target_f <= freqs[-1]:
            target_idx = np.argmin(np.abs(freqs - target_f))
            target_amp = amp_mean[target_idx]
            target_std = amp_std[target_idx]
            target_snr = target_amp / target_std if target_std > 0 else float('inf')
            
            print(f"\n目标频率 {target_f} Hz 分析:")
            print(f"平均幅度: {target_amp:.6g} µV")
            print(f"标准差: {target_std:.6g} µV")
            print(f"信噪比: {target_snr:.2f}")
        
        # 谐波分析
        print(f"\n谐波检测结果:")
        for k in range(1, 5):
            fk = k * target_f
            if k in found:
                fx, fy = found[k]
                print(f"{k}f: 频率≈{fx:.3f} Hz, 幅度={fy:.6g} µV")
            else:
                print(f"{k}f: 未检测到明显峰值")

        print(f'完成：{base_name} 的5-15秒窗口化幅度谱已保存到 {out_dir}')
        return True
        
    except Exception as e:
        print(f"处理文件 {vhdr_path} 时出错: {str(e)}")
        return False

def main():
    # ========= 用户可修改区域 =========
    data_dir = "D:\\code\\Python\\library\\data\\2025.9.4\\AR_three"  # 数据目录
    out_dir = os.path.join(os.path.dirname(__file__), 'D:\\code\\Python\\library\\data\\2025.9.4\\result_AR_three')  # 输出目录
    # 窗口参数
    window_length = 4.0  # 窗口长度（秒）
    overlap = 0.5        # 重叠比例
    # =================================

    ensure_dir(out_dir)
    
    # 查找所有vhdr文件
    vhdr_pattern = os.path.join(data_dir, "*.vhdr")
    vhdr_files = glob.glob(vhdr_pattern)
    
    if not vhdr_files:
        print(f"在目录 {data_dir} 中未找到任何.vhdr文件")
        return
    
    print(f"找到 {len(vhdr_files)} 个vhdr文件，开始批量处理...")
    
    # 按文件名排序（确保处理顺序一致）
    vhdr_files.sort()
    
    success_count = 0
    total_count = len(vhdr_files)
    
    for i, vhdr_path in enumerate(vhdr_files, 1):
        print(f"\n进度: {i}/{total_count}")
        if process_single_file(vhdr_path, out_dir, window_length, overlap):
            success_count += 1
    
    print(f"\n批量处理完成！")
    print(f"成功处理: {success_count}/{total_count} 个文件")
    print(f"输出目录: {out_dir}")

if __name__ == '__main__':
    main()
